"""
基于RSRS指标的多标的策略
使用分钟K线，多品种RSRS策略
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask
from sklearn.linear_model import LinearRegression

class MutilRsrsStrategy:
    strategy_name = "MutilRsrsStrategy"
    """
    多标的RSRS策略
    
    RSRS (阻力支撑相对强度) 指标原理:
    1. 计算一段时间内最高价对最低价的线性回归斜率
    2. 对斜率进行标准化处理
    3. 斜率越大表示上涨趋势越强，斜率越小表示下跌趋势越强
    
    策略逻辑:
    - RSRS > 上阈值: 做多
    - RSRS < 下阈值: 做空
    - 其他情况: 平仓或持有
    """
    
    def __init__(self, api, symbol_info, market_period):
        self.api = api
        self.all_info = {}
        i_symbols = []
        
        # 初始化每个品种的信息
        for symbol, info in symbol_info.items():
            self.all_info[symbol] = info
            i_symbol = f"KQ.i@{info['exchange']}.{symbol}"
            self.all_info[symbol]["i_symbol"] = i_symbol
            i_symbols.append(i_symbol)
            self.all_info[symbol]["m_symbol"] = f"KQ.m@{info['exchange']}.{symbol}"
            self.all_info[symbol]["rsrs_slopes"] = []  # RSRS斜率序列
            self.all_info[symbol]["last_trade_time"] = 0  # 上次交易时间
            self.all_info[symbol]["entry_price"] = 0.0  # 入场价格

        # 数据订阅 - 使用指定周期K线
        self.kline = api.get_kline_serial(i_symbols, market_period * 60, data_length=800)
        self.market_period = market_period
        
        # 交易状态
        self.run_count = 0
        
    def calculate_rsrs(self, symbol, kline) -> float:
        """
        计算RSRS指标
        
        Args:
            symbol: 品种代码
            kline: K线数据
            
        Returns:
            标准化后的RSRS值
        """
        try:
            rsrs_period = self.all_info[symbol]["rsrs_period"]
            rsrs_std_period = self.all_info[symbol]["rsrs_std_period"]
            
            if len(kline) < rsrs_std_period:
                return 0.0
            
            # 获取最近的数据
            high_prices = kline.high.values
            low_prices = kline.low.values
            
            rsrs_slopes = self.all_info[symbol]["rsrs_slopes"]
            
            # 首次计算，从足够的历史数据开始计算
            if len(rsrs_slopes) < rsrs_std_period:
                for i in range(rsrs_period, len(high_prices)-1):
                    # 获取计算窗口内的高低价
                    window_high = high_prices[i-rsrs_period:i]
                    window_low = low_prices[i-rsrs_period:i]
                    
                    # 线性回归计算斜率
                    X = window_low.reshape(-1, 1)  # 最低价作为自变量
                    y = window_high  # 最高价作为因变量
                    
                    reg = LinearRegression().fit(X, y)
                    slope = reg.coef_[0]
                    self.all_info[symbol]["rsrs_slopes"].append(slope)
            # 后续更新计算
            else:
                i = len(high_prices) - 1
                window_high = high_prices[i-rsrs_period:i]
                window_low = low_prices[i-rsrs_period:i]
                
                # 线性回归计算斜率
                X = window_low.reshape(-1, 1)  # 最低价作为自变量
                y = window_high  # 最高价作为因变量
                reg = LinearRegression().fit(X, y)
                slope = reg.coef_[0]
                self.all_info[symbol]["rsrs_slopes"].append(slope)
            
            # 获取最新的RSRS斜率
            current_slope = self.all_info[symbol]["rsrs_slopes"][-1]
            
            # 标准化处理
            recent_slopes = self.all_info[symbol]["rsrs_slopes"][-rsrs_std_period:]
            mean_slope = np.mean(recent_slopes)
            std_slope = np.std(recent_slopes)
            
            if std_slope > 0:
                rsrs_score = (current_slope - mean_slope) / std_slope
            else:
                rsrs_score = 0.0
            
            return rsrs_score
            
        except Exception as e:
            print(f"RSRS计算失败 {symbol}: {e}")
            return 0.0
    
    def calculate_atr(self, kline, period: int = 14) -> float:
        """计算ATR用于止损止盈"""
        try:
            if len(kline) < period + 1:
                return kline.close.iloc[-1] * 0.02  # 默认2%
            
            high = kline.high.iloc[-period:]
            low = kline.low.iloc[-period:]
            close = kline.close.iloc[-period-1:-1]  # 前一日收盘价
            
            tr1 = high - low
            tr2 = abs(high - close.values)
            tr3 = abs(low - close.values)
            
            tr = np.maximum(tr1, np.maximum(tr2, tr3))
            atr = np.mean(tr)
            
            return atr
            
        except Exception as e:
            print(f"ATR计算失败: {e}")
            return kline.close.iloc[-1] * 0.02
    
    def on_bar(self):
        """主循环"""
        bar_count = 0
        while True:
            self.api.wait_update()
            if self.api.is_changing(self.kline):
                bar_count += 1
                current_time = self.kline.datetime.iloc[-1]
                current_price = self.kline.close.iloc[-1]
                volume = self.kline.volume.iloc[-1]

                # 每50根K线打印一次进度
                if bar_count % 50 == 0:
                    print(f"处理第{bar_count}根{self.market_period}分钟K线: 时间={current_time}, 价格={current_price:.1f}")

                for symbol in self.all_info.keys():
                    if volume != 0:
                        self.run(symbol)
                    if volume == 0:
                        self.kline = self.kline.iloc[:-1]
                        self.run(symbol)
    
    def run(self, symbol):
        """策略执行逻辑"""
        try:
            m_symbol = self.all_info[symbol]["m_symbol"]
            symbol_info = self.api.get_quote(m_symbol)
            main_symbol = symbol_info.underlying_symbol
            target_pos = TargetPosTask(self.api, main_symbol)
            
            # 计算每个品种的仓位大小
            symbol_value = symbol_info.volume_multiple * symbol_info.last_price
            position_size = int(self.api.get_account().balance / symbol_value / len(self.all_info.keys()))
            if position_size < 1:
                position_size = 1
            
            # 平掉非当前主力的持仓
            positions = self.api.get_position()
            for sym, pos in positions.items():
                if symbol in sym:
                    if sym != main_symbol and pos.pos != 0:
                        TargetPosTask(self.api, sym).set_target_volume(0)
            
            main_pos = positions.get(main_symbol, None)
            current_pos = main_pos.pos if main_pos else 0
            current_price = symbol_info.last_price
            current_time = self.kline.datetime.iloc[-1]

            # 获取对应品种的K线数据
            for i in range(len(self.all_info.keys())):
                if symbol in self.kline["symbol" if i==0 else "symbol"+str(i)].iloc[-1]:
                    add_str = "" if i==0 else str(i)
                    cols = ["datetime","id"+add_str,"open"+add_str,"high"+add_str,"low"+add_str,
                                "close"+add_str,"volume"+add_str,"open_oi"+add_str,"duration"]
                    kline = self.kline[cols].rename(
                        columns = {"id"+add_str:"id",
                                    "open"+add_str:"open",
                                    "high"+add_str:"high",
                                    "low"+add_str:"low",
                                    "close"+add_str:"close",
                                    "volume"+add_str:"volume",
                                    "open_oi"+add_str:"open_oi"}
                    )
                    break

            # 检查数据是否足够
            if len(kline) < self.all_info[symbol]["rsrs_std_period"]:
                return
            
            # 控制交易频率
            if current_time - self.all_info[symbol]["last_trade_time"] < self.market_period * 60 * 1000:
                return
            
            # 计算RSRS指标
            rsrs_score = self.calculate_rsrs(symbol, kline)
            
            # 调试信息：每20次运行打印一次
            self.run_count += 1
            if self.run_count < self.all_info[symbol]["start_period"]:
                return

            if self.run_count % 20 == 0:
                print(f"{symbol} RSRS策略运行{self.run_count}次: RSRS={rsrs_score:.3f}, 持仓={current_pos}, 价格={current_price:.1f}")
            
            long_threshold = self.all_info[symbol]["long_threshold"]
            short_threshold = self.all_info[symbol]["short_threshold"]
            close_threshold = self.all_info[symbol]["close_threshold"]
            
            if current_pos == 0:  # 无持仓
                if rsrs_score > long_threshold:  # 做多信号
                    # 做多
                    self.all_info[symbol]["entry_price"] = current_price
                    target_pos.set_target_volume(position_size)
                    self.all_info[symbol]["last_trade_time"] = current_time
                    
                    print(f"🟢 {symbol} RSRS做多: 指标={rsrs_score:.3f}, 价格={current_price:.1f}")
                
                elif rsrs_score < short_threshold:  # 做空信号
                    # 做空
                    self.all_info[symbol]["entry_price"] = current_price
                    target_pos.set_target_volume(-position_size)
                    self.all_info[symbol]["last_trade_time"] = current_time
                    
                    print(f"🔴 {symbol} RSRS做空: 指标={rsrs_score:.3f}, 价格={current_price:.1f}")
            
            else:  # 有持仓
                should_close = False
                close_reason = ""
                
                if current_pos > 0:  # 多头持仓
                    if rsrs_score < -close_threshold:  # RSRS信号反转
                        should_close = True
                        close_reason = f"RSRS信号反转({rsrs_score:.3f})"
                
                elif current_pos < 0:  # 空头持仓
                    if rsrs_score > close_threshold:  # RSRS信号反转
                        should_close = True
                        close_reason = f"RSRS信号反转({rsrs_score:.3f})"
                
                if should_close:
                    # 平仓
                    target_pos.set_target_volume(0)
                    self.all_info[symbol]["last_trade_time"] = current_time
                    
                    print(f"⚪ {symbol} {close_reason}: 价格={current_price:.1f}")
                    
        except Exception as e:
            print(f"{symbol} 策略执行失败: {e}")